#include <iostream>
#include <mutex>
#include <shared_mutex>
#include <vector>
#include <thread>
#include <condition_variable>
#include <queue>
#include <unordered_map>
#include <set>
#include <algorithm>
#include <atomic>


const int N = 10;
std::vector<std::vector<int>> test_idx_;
std::vector<int> test_result_;


//数组中的元素， 每一个元素都会有一把读写锁
struct elem {
    int val_;                      //储存的元素
};

//数组：所有数字输出化为1， 方便调试
struct ElemS {
    elem elems_[N];

    ElemS() {
        for (int i = 0; i < N; i++) {
            elems_[i].val_ = 1;
        }
    }

    int& operator[](int i) {
        return elems_[i].val_;
    }
};

//设计锁管理器
struct MutexManager {
    //全局的事务id， 每执行一次循环， 就+1
    int txn_id_ = 0;

    //每个锁的执行事务： 锁id -> 事务id
    std::unordered_map<int, int> mid_w_tid_; 
    std::unordered_map<int, std::set<int>> mid_r_tid_; 

    //每个锁以及对应的等待队列: 锁id -> 等待队列
    std::unordered_map<int, std::queue<int>> mid_q_;

    //锁存器
    std::mutex latch_;
    std::condition_variable cv_;


    //计数器， 方便调试
    std::atomic<int> timestramp_ = 0;
};


//工作器， Request为1w次请求
struct Worker {
    //释放idx->txn_id的读锁
    auto ReleaseReadMutex(MutexManager &mm, int idx, int txn_id) {
        mm.mid_r_tid_[idx].erase(txn_id);
        if (mm.mid_r_tid_[idx].size() == 0) {
            mm.mid_r_tid_.erase(idx);
        }
    }
    
    //释放idx位置的写锁
    auto ReleaseWriteMutex(MutexManager &mm, int idx) {
        mm.mid_w_tid_.erase(idx);
    }

    //释放全部的锁， 用于一次操作完成后释放所有锁
    auto ReleaseAllMutex(ElemS &s, MutexManager &mm, int txn_id, int i, int j) {
        std::unique_lock<std::mutex> l(mm.latch_);        
        test_idx_.push_back({j, i, (i + 1) % N, (i + 2) % N});
        test_result_.push_back(s[j]);

        //释放s[j]
        ReleaseWriteMutex(mm, j);
        //释放s[i]
        if (j != i) {
            ReleaseReadMutex(mm, i, txn_id);
        }
        //释放s[i + 1]
        if (j != (i + 1) % N) {
            ReleaseReadMutex(mm, (i + 1) % N, txn_id);
        }
        //释放s[i + 2]
        if (j != (i + 2) % N) {
            ReleaseReadMutex(mm, (i + 2) % N, txn_id);
        }

        mm.cv_.notify_all();
    }

    //获得指定位置的锁
    auto GetMutex(MutexManager &mm, std::unique_lock<std::mutex> &l, int txn_id, 
                    int j, int i, std::vector<int> prev_idxs) -> bool {
        
        if (j != i) {
            //判断是否有写锁了， 如果有进入if
            if (mm.mid_w_tid_.count(i) != 0) {
                //等待-死亡策略， 如果我更年轻， 那么我应该被中止
                if (txn_id > mm.mid_w_tid_[i]) {
                    // prev_ids保存之前获取的锁的位置，然后把这些锁释放
                    for (int k = 0; k < prev_idxs.size(); k++) {
                        if (prev_idxs[k] == j) {
                            mm.mid_w_tid_.erase(prev_idxs[k]);
                        } else {
                            mm.mid_r_tid_[prev_idxs[k]].erase(txn_id);
                            if (mm.mid_r_tid_[prev_idxs[k]].size() == 0) { mm.mid_r_tid_.erase(prev_idxs[k]); }
                        }
                    }

                    mm.cv_.notify_all();
                    return false;
                }

                mm.mid_q_[i].push(txn_id);
                
                mm.cv_.wait(l, [&](){
                    return mm.mid_w_tid_.count(i) == 0 
                    && mm.mid_q_[i].front() == txn_id;
                });  

                mm.mid_q_[i].pop();

                
                mm.mid_r_tid_[i].insert(txn_id);
            } else {
                mm.mid_r_tid_[i].insert(txn_id);
            }

        } else {

            //获取写锁， 此时查看是否已经获取过写锁， 如果已经获取过， 那么返回不需要重复获取
            for (int k = 0; k < prev_idxs.size(); k++) {
                if (j == prev_idxs[k]) {
                    return true;
                }
            }

            // 获取s[j]的写锁, 如果被占用，则等待， 并被加入等待队列
            if (mm.mid_r_tid_.count(j) || mm.mid_w_tid_.count(j)) {

                // 等待-死亡策略
                // 如果我更年轻， 那么我应该被中止. 对于写锁， 如果我比写锁中最年长的年轻， 我应该中止
                if (mm.mid_r_tid_.count(j)) {
                    //先检查读锁冲突
                    if (txn_id > (*mm.mid_r_tid_[j].begin())) {
                        //当前为获取写锁， 所以只需要释放读锁：
                        for (int k = 0; k < prev_idxs.size(); k++) {
                            mm.mid_r_tid_[prev_idxs[k]].erase(txn_id);
                            if(mm.mid_r_tid_[prev_idxs[k]].size() == 0) { mm.mid_r_tid_.erase(prev_idxs[k]); }
                        }

                        mm.cv_.notify_all();
                        return false;
                    }
                } else if (mm.mid_w_tid_.count(j)) {
                    if (txn_id > mm.mid_w_tid_[j]) {
                            //当前为获取写锁， 所以只需要释放读锁：
                            for (int k = 0; k < prev_idxs.size(); k++) {
                                mm.mid_r_tid_[prev_idxs[k]].erase(txn_id);
                                if(mm.mid_r_tid_[prev_idxs[k]].size() == 0) { mm.mid_r_tid_.erase(prev_idxs[k]); }
                            }
                        }

                        mm.cv_.notify_all();
                        return false;
                    }

                // 加入等待队列
                mm.mid_q_[j].push(txn_id);

                // 锁住， 等待被唤醒， 唤醒指定条件为该位置读锁和写锁都没被占用， 并且等待队头为自己
                mm.cv_.wait(l, [&]() {
                    return ((mm.mid_r_tid_.count(j) == 0) || mm.mid_r_tid_[j].size() == 0)
                    && mm.mid_w_tid_.count(j) == 0 
                    && mm.mid_q_[j].front() == txn_id;
                });

                // 删除等待队列中的自己，并且将自己的写标记写入锁管理器
                mm.mid_q_[j].pop();

                mm.mid_w_tid_[j] = txn_id;
            } else {// s[j]锁都没有被占用
                mm.mid_w_tid_[j] = txn_id;
            }
        }
        return true;
    }

    //v1版本： 2PL + 死锁预防(等待-死亡)
    auto GetAllMutex_v1(MutexManager &mm, int txn_id, std::vector<int> idxs) -> bool {
                    //   j         i      (i+1)%N   (i+2)%N
        int j = idxs[0];
        std::vector<int> already_getidxs;
        
        for (int i = 0; i < idxs.size(); i++) {
            std::unique_lock<std::mutex> l(mm.latch_);
            
            if (!GetMutex(mm, l, txn_id, j, idxs[i], already_getidxs)) {
                return false;
            }
            already_getidxs.push_back(idxs[i]);
        }

        return true;   
    }


    //v2版本 : 2PL + 死锁预防 + 锁顺序获取
    auto GetAllMutex_v2(MutexManager &mm, int txn_id, std::vector<int> idxs) -> bool {

        int j = idxs[0];

        //v1和v2的唯一不同点在于这里排了序
        sort(idxs.begin(), idxs.end());
        //已经获取到的锁
        std::vector<int> alreaddy_getidxs;

        for (int i = 0; i < idxs.size(); i++) {
            std::unique_lock<std::mutex> l(mm.latch_);
            if (!GetMutex(mm, l, txn_id, j, idxs[i], alreaddy_getidxs)) {
                return false;
            }
            alreaddy_getidxs.push_back(idxs[i]);
        }
        return true;
    }

    auto Request(ElemS &s,  MutexManager &mm, int i) {
        srand(time(0) + i);
        for (int k = 0; k < 10000; k++) {
            // 1.获取随机数字， 模拟事务并获取事务id。
            int i = rand() % N;
            int j = rand() % N;
            std::vector<int> idxs{j, i, (i + 1) % N, (i + 2) % N}; //idxs为要处理的四个位置

            // 2.获取执行id
            int txn_id = ++mm.txn_id_;

            // 3.获取锁
            if (!GetAllMutex_v1(mm, txn_id, idxs)) {
                continue; // 失败则执行下一次操作
            }
            
            // 4.进行计算
            s[j] = s[i] + s[(i + 1) % N] + s[(i + 2) % N];
            std::cout << s[j] << " = " << s[i] << " + " 
            << s[(i + 1) % N] << " + " << s[(i + 2) % N] << "  time:" << mm.timestramp_++ << std::endl;

            // 5.释放锁
            ReleaseAllMutex(s, mm, txn_id, i, j);
        }
    }
};


struct Workers {

    //初始化， 生成多个工作线程。 
    auto Init(int count = 2) {
        count_ = count;
        threads_.resize(count_);
        workers_.resize(count_);
    }

    auto Run(ElemS &s, MutexManager &mm) {
        //后台执行所有工作器
        for (int i = 0; i < count_; i++) {
            threads_[i] = std::thread(&Worker::Request, &workers_[i], std::ref(s), std::ref(mm), i);
        }

        for (int i = 0; i    < count_; i++) {
            threads_[i].join();
        }
    }

    int count_;
    std::vector<std::thread> threads_;
    std::vector<Worker> workers_;
};


int main() {
    ElemS s;            //创建数组
    auto test_s = s;    //拷贝数组用于测试
    
    MutexManager mm;

    Workers ws;
    
    ws.Init(10);
    
    ws.Run(s, mm);
    std::cout << "finish" << std::endl;

    for (int i = 0; i < test_idx_.size(); i++) {
        auto result = test_result_[i];
        int idx1 = test_idx_[i][0];
        int idx2 = test_idx_[i][1];
        int idx3 = test_idx_[i][2];
        int idx4 = test_idx_[i][3];


        test_s[idx1] = test_s[idx2] + test_s[idx3] + test_s[idx4]; 
        if (test_s[idx1] != result) {
            std::cout << "test fialed!" << std::endl;
            return 0;
        }
    }

    std::cout << "test success!" << std::endl;

    return 0;
}

